from flask import Flask, Response, stream_with_context
import cv2
from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import get_graph_path, model_wh
from tf_pose import common

app = Flask(__name__)


def generate_frames():
    # 初始化姿态估计模型
    e = TfPoseEstimator(get_graph_path('cmu'), target_size=(432, 368))

    video_capture = cv2.VideoCapture(0)  # 从摄像头读取视频

    while True:
        ret, frame = video_capture.read()
        if not ret:
            break

            # 姿态估计处理
        humans = e.inference(frame, resize_to_default=True, upsample_size=4.0)
        frame = TfPoseEstimator.draw_humans(frame, humans, imgcopy=False)

        # 绘制坐标和节点名称
        for human in humans:
            for body_part in human.body_parts.values():
                if body_part.score > 0.1:  # 过滤掉得分较低的点
                    # 计算并绘制坐标点
                    center = (int(body_part.x * frame.shape[1]), int(body_part.y * frame.shape[0]))
                    cv2.circle(frame, center, 3, common.CocoColors[body_part.part_idx], thickness=3, lineType=8,
                               shift=0)

                    # 绘制节点名称
                    text = f"{common.CocoPart(body_part.part_idx).name}: ({center[0]}, {center[1]})"
                    cv2.putText(frame, text, (center[0], center[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                                (255, 255, 255), 2, cv2.LINE_AA)

                    # 将图像编码为 JPEG 格式
        _, buffer = cv2.imencode('.jpg', frame)
        frame = buffer.tobytes()

        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')


@app.route('/tfposecamera')
def tfpose_camera():
    return Response(stream_with_context(generate_frames()),
                    mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=10001, debug=True)